from typing import Literal
import asyncio
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool, StructuredTool, BaseTool, ToolException
from pydantic import BaseModel, Field

from langChain.config import model


# 创建工具的3种方法 ：
#     1. 使用@tool装饰器，最简单的方式
#     2. 使用StructuredTool.from_function 类方法 --这类似于 @too1 装饰器，但允许更多配置和同步和异步实现的规范
#     3. 通过子类化BaseTool--这是最灵活的方法，它提供了最大程度的控制，但需要更多的工作量和代码。
#
# 多数用例 ，使用 @too1或StructuredTool.from_function足够了。
# 如果工具具有精心选择的名称、描述和 JSON 模式，BaseTool子类 模型的性能会更好。


# 第一种，@tool装饰器
@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers."""
    return a * b


# # 让我们检查与该工具关联的一些属性。
# print('@tool装饰器名称 ：', multiply.name)
# print('@tool装饰器描述 ：', multiply.description)
# print('@tool装饰器参数 ：', multiply.args)
#
# 异步工具
# @tool(name_or_callable="第二个工具",description="工具描述",)
# async def amultiply(a: int, b: int) -> int:
#     """Multiply two numbers."""
#     return a * b
#
#
# # 让我们检查与该工具关联的一些属性。
# print('@tool异步装饰器名称 ：', amultiply.name)
# print('@tool异步装饰器描述 ：', amultiply.description)
# print('@tool异步装饰器参数 ：', amultiply.args)
#
# # 自定义参数类型
# class CalculatorInput(BaseModel):
#     a: int = Field(description="first number")
#     b: int = Field(description="second number")
# @tool(args_schema=CalculatorInput,return_direct=True)
# def multiply(a: int, b: int) -> int:
#     """Multiply two numbers."""
#     return a * b
# print('@tool装饰器 【自定义了参数】 名称 ：', multiply.name)
# print('@tool装饰器 【自定义了参数】 描述 ：', multiply.description)
# print('@tool装饰器 【自定义了参数】 参数 ：', multiply.args)
# print('@tool装饰器 【自定义了参数】 参数类型 ：', multiply.args_schema)
# print('@tool装饰器 【自定义了参数】 return_direct ：', multiply.return_direct)
# print('===============================================')

# # 第二种，StructuredTool.from_function
#
# def multiply(a: int, b: int) -> int:
#     """两个数相乘."""
#     return a * b
#
#
# async def amultiply(a: int, b: int) -> int:
#     """Multiply two numbers."""
#     return a * b
#
#
# async def main():
#     # func 参数：指定一个同步函数。当你在同步上下文中调用工具时，它会使用这个同步函数来执行操作。
#     # oroutine 参数：指定一个异步函数。当你在异步上下文中调用工具时，它会使用这个异步函数来执行操作。
#     calculator = StructuredTool.from_function(func=multiply, coroutine=amultiply)
#     # invoke是同步调用
#     print(calculator.invoke({"a": 2, "b": 3}))
#     # ainvoke是异步调用
#     print(await calculator.ainvoke({"a": 2, "b": 5}))
#
#
# # 运行异步主函数
# asyncio.run(main())


# # StructuredTool自定义配置参数
# class CalculatorInput(BaseModel):
#     a: int = Field(description="first number")
#     b: int = Field(description="second number")
#
#
# def multiply(a: int, b: int) -> int:
#     """Multiply two numbers."""
#     return a * b
#
#
# # 创建一个异步包装器函数
# async def async_addition(a: int, b: int) -> int:
#     """Multiply two numbers."""
#     return a + b
#
#
# async def main():
#     calculator = StructuredTool.from_function(
#         func=multiply,
#         coroutine=async_addition,
#         name="Calculator",
#         description="multiply numbers",
#         args_schema=CalculatorInput,
#         return_direct=True,
#     )
#     print(calculator.invoke({"a": 2, "b": 3}))
#     print(await calculator.ainvoke({"a": 2, "b": 5}))
#     print(calculator.name)
#     print(calculator.description)
#     print(calculator.args)
#
#
# # 运行异步主函数
# asyncio.run(main())
# print("===============================================")

# # 第三种，BaseTool子类
# class Calculator(BaseTool):
#     name: str = "BaseTool子类"
#     description: str = "通过继承 BaseTool 来创建工具的"
#
#     # 同步处理
#     def _run(self, a: int, b: int) -> int:
#         """Multiply two numbers."""
#         return a * b
#
#     # 异步处理
#     def _arun(self, a: int, b: int) -> int:
#         """Multiply two numbers."""
#         return a * b
#
#
# calculator = Calculator()
# print(calculator.invoke({"a": 2, "b": 3}))
# print(calculator.ainvoke({"a": 2, "b": 5}))
# print(calculator.name)
# print(calculator.description)
# print(calculator.args)
# print("===============================================")


# tool 的错误处理机制
# 一个简单的策略是在工具内部抛出ToolException ，然后使用 handle_tool_error 指定一个错误处理程序。默认handle_tool_error = False 。
# handle_tool_eror 设置为True 、字符串值、函数，函数接受一个参数，即错误对象。如果是函数，该函数应该以 ToolException 作为参数，并返回一个值。


def get_weather(city: str) -> int:
    """获取给定城市的天气。"""
    raise ToolException(f"错误：没有名为{city}的城市。")


# # 默认情况，直接抛出异常
# get_weather_tool = StructuredTool.from_function(
#     func=get_weather,
#     # 默认情况下，如果函数抛出ToolException，则将ToolException的message作为响应。
# )
# response = get_weather_tool.invoke({"city": "foobar"})
# print('错误的默认处理方式【直接抛出异常】 :', response)

# # handle_tool_error=True，则将返回ToolException异常文本
# get_weather_tool = StructuredTool.from_function(
#     func=get_weather,
#     # 如果设置为True，则将返回ToolException异常文本，False将会抛出ToolException
#     handle_tool_error=True,
# )
# response = get_weather_tool.invoke({"city": "foobar"})
# print('handle_tool_error=True 的处理方式 【返回异常文本】:', response)

# # handle_tool_error设置str,则将返回指定的文本
# get_weather_tool = StructuredTool.from_function(
#     func=get_weather,
#     handle_tool_error="没找到这个城市",
# )
#
# response = get_weather_tool.invoke({"city": "foobar"})
# print('handle_tool_error==str 的处理方式 【返回设置的文本】:', response)

# # handle_tool_error设置fun,则将返回函数的返回值
# def _handle_error(error: ToolException) -> str:
#     return f"工具执行期间发生以下错误：`{error.args[0]}`"
#
# get_weather_tool = StructuredTool.from_function(
#     func=get_weather,
#     handle_tool_error=_handle_error,
# )
#
# response = get_weather_tool.invoke({"city": "foobar"})
# print('handle_tool_error==fun 的处理方式 【返回函数的返回值】:',response)
# print("===============================================")

# 完整的例子
@tool
def weather_tool(weather: Literal["晴朗的", "多云的", "多雨的", "下雪的"]) -> None:
    """Describe the weather"""
    pass


# 模型绑定工具
model_with_tools = model.bind_tools([weather_tool])

image_url_1 = "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
image_url_2 = "https://upload.wikimedia.org/wikipedia/commons/thumb/9/92/Morning_in_China_Snow_Town.jpg/1280px-Morning_in_China_Snow_Town.jpg"

message = HumanMessage(
    content=[
        {"type": "text", "text": "用中文描述两张图片中的天气"},
        {"type": "image_url", "image_url": {"url": image_url_1}},
        {"type": "image_url", "image_url": {"url": image_url_2}},
    ],
)
response = model_with_tools.invoke([message])
print(response.tool_calls)
